/*=========================================================================

Library:   TubeTK

Copyright 2010 Kitware Inc. 28 Corporate Drive,
Clifton Park, NY, 12065, USA.

All rights reserved.

Licensed under the Apache License, Version 2.0 ( the "License" );
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

=========================================================================*/

#include "tubeMessage.h"

#include <itkImageFileReader.h>
#include <itkImageFileWriter.h>

// TubeTKITK includes
#include "tubeConvertSpatialGraphToImage.h"

#include "ConvertSpatialGraphToImageCLP.h"

template< class TPixel, unsigned int VDimension >
int DoIt( int argc, char * argv[] );

// Must follow include of "...CLP.h" and forward declaration of int DoIt( ... ).
#include "../CLI/tubeCLIHelperFunctions.h"

int main( int argc, char * argv[] )
{
  PARSE_ARGS;
  return tube::ParseArgsAndCallDoIt( inCentroidsFile, argc, argv );
}

template< class TPixel, unsigned int VDimension >
int DoIt( int argc, char * argv[] )
{
  PARSE_ARGS;

  std::stringstream logMsg;

  typedef itk::Image< float, VDimension >       ImageType;
  typedef itk::Image< TPixel, VDimension >      CVTImageType;
  typedef itk::ImageFileReader< CVTImageType >  ImageReaderType;

  typedef tube::ConvertSpatialGraphToImage< CVTImageType, ImageType >
    ConvertSpatialGraphToImageFilterType;
  typename ConvertSpatialGraphToImageFilterType::Pointer filter =
    ConvertSpatialGraphToImageFilterType::New();

  typename ImageReaderType::Pointer imageReader = ImageReaderType::New();

  // Read image file generated by CVT
  imageReader->SetFileName( inCentroidsFile.c_str() );
  imageReader->Update();

  filter->SetInput( imageReader->GetOutput() );

  double temp;
  int numberOfCentroids_AdjMat;
  int numberOfCentroids;
  std::ifstream readStream;

  // Read full connectivity information ( Adjacency Matrix )
  std::string matrixFilename = inGraphFile + ".mat";

  logMsg.str( "" );
  logMsg << "Reading file: " << matrixFilename;
  tube::InfoMessage( logMsg.str() );

  readStream.open( matrixFilename.c_str(), std::ios::in );
  readStream >> numberOfCentroids_AdjMat;
  readStream.get();

  logMsg.str( "" );
  logMsg << "Number of centroids: " << numberOfCentroids_AdjMat;
  tube::InfoMessage( logMsg.str() );

  vnl_matrix< double > adjacencyMatrix( numberOfCentroids_AdjMat,
    numberOfCentroids_AdjMat );
  adjacencyMatrix.fill( 0 );
  vnl_vector< double > branchnessVector( numberOfCentroids_AdjMat );
  branchnessVector.fill( 0 );
  vnl_vector< double > radiusVector( numberOfCentroids_AdjMat );
  radiusVector.fill( 0 );
  vnl_vector< double > centralityVector( numberOfCentroids_AdjMat );
  centralityVector.fill( 0 );
  for( int i = 0; i < numberOfCentroids_AdjMat; i++ )
    {
    for( int j = 0; j < numberOfCentroids_AdjMat; j++ )
      {
      readStream >> temp;
      readStream.get();
      adjacencyMatrix[i][j] = temp;
      }
    }
  filter->SetAdjacencyMatrix( adjacencyMatrix );
  readStream.close();

  // Read in BRANCH file
  std::string branchFilename = inGraphFile + ".brc";

  logMsg.str( "" );
  logMsg << "Reading ( branch ) file: " << branchFilename;
  tube::InfoMessage( logMsg.str() );

  readStream.open( branchFilename.c_str(), std::ios::in );
  readStream >> numberOfCentroids;
  readStream.get();
  if( numberOfCentroids != numberOfCentroids_AdjMat )
    {
    tube::ErrorMessage( "Error: fileList's #Centroids != branch #Centroids" );
    return EXIT_FAILURE;
    }
  for( int i = 0; i < numberOfCentroids; i++ )
    {
    readStream >> temp;
    readStream.get();
    branchnessVector[i] = temp;
    }
  filter->SetBranchnessVector( branchnessVector );
  readStream.close();

  // Read in ROOT file
  std::string rootFilename = inGraphFile + ".rot";

  logMsg.str( "" );
  logMsg << "Reading ( root ) file: " << rootFilename;
  tube::InfoMessage( logMsg.str() );

  readStream.open( rootFilename.c_str(), std::ios::in );
  readStream >> numberOfCentroids;
  readStream.get();
  if( numberOfCentroids != numberOfCentroids_AdjMat )
    {
    tube::ErrorMessage( "Error: fileList's #Centroids != root #Centroids" );
    return EXIT_FAILURE;
    }
  for( int i = 0; i < numberOfCentroids; i++ )
    {
    readStream >> temp;
    readStream.get();
    radiusVector[i] = temp;
    }
  filter->SetRadiusVector( radiusVector );
  // readStream.close();
  // Read in CENTRALITY file
  // std::string centralityFilename = inGraphFile + ".cnt";

  // logMsg.str( "" );
  // logMsg << "Reading ( centrality ) file: " << centralityFilename;
  // tube::InfoMessage( logMsg.str() );

  // readStream.open( centralityFilename.c_str(), std::ios::in );
  // readStream >> numberOfCentroids;
  // readStream.get();
  // if( numberOfCentroids != numberOfCentroids_AdjMat )
  //   {
  //   tube::ErrorMessage( "Error: fileList's #Centroids != centrality
  //     #Centroids" );
  //   return EXIT_FAILURE;
  //   }
  for( int i = 0; i < numberOfCentroids; i++ )
    {
    readStream >> temp;
    readStream.get();
    centralityVector[i] = temp;
    }
  filter->SetCentralityVector( centralityVector );
  readStream.close();

  filter->Update();

  typedef itk::ImageFileWriter< ImageType > ImageWriterType;
  typename ImageWriterType::Pointer imageWriter;
  std::string filename;

  filename = outImageFile + ".mat.mha";
  imageWriter = ImageWriterType::New();
  imageWriter->SetFileName( filename.c_str() );
  imageWriter->SetInput( filter->GetAdjacencyMatrixImage() );
  imageWriter->Write();

  filename = outImageFile + ".brc.mha";
  imageWriter = ImageWriterType::New();
  imageWriter->SetFileName( filename.c_str() );
  imageWriter->SetInput( filter->GetBranchnessImage() );
  imageWriter->Write();

  filename = outImageFile + ".rot.mha";
  imageWriter = ImageWriterType::New();
  imageWriter->SetFileName( filename.c_str() );
  imageWriter->SetInput( filter->GetRadiusImage() );
  imageWriter->Write();

  filename = outImageFile + ".cnt.mha";
  imageWriter = ImageWriterType::New();
  imageWriter->SetFileName( filename.c_str() );
  imageWriter->SetInput( filter->GetCentralityImage() );
  imageWriter->Write();

  return EXIT_SUCCESS;
}
